import torch
from torch import nn
from torch.nn import functional as F
# https://zh-v2.d2l.ai/chapter_convolutional-modern/resnet.html#id4

class Residual(nn.Module):  #@save
    def __init__(self, input_channels, num_channels,
                 use_1x1conv=False, strides=1):
        super().__init__()
        self.conv1 = nn.Conv2d(input_channels, num_channels,
                               kernel_size=3, padding=1, stride=strides)
        self.conv2 = nn.Conv2d(num_channels, num_channels,
                               kernel_size=3, padding=1)
        # 当 use_1x1conv = True 时，添加通过1*1卷积调整通道和分辨率。
        # 当 use_1x1conv = False 时，应用ReLU非线性函数之前，将输入添加到输出。
        if use_1x1conv:
            self.conv3 = nn.Conv2d(input_channels, num_channels,
                                   kernel_size=1, stride=strides)
        else:
            self.conv3 = None
        self.bn1 = nn.BatchNorm2d(num_channels)
        self.bn2 = nn.BatchNorm2d(num_channels)

    def forward(self, X):
        # 每个残差块有2个卷积层（不包括恒等映射的1*1卷积层）
        Y = F.relu(self.bn1(self.conv1(X)))
        Y = self.bn2(self.conv2(Y))
        if self.conv3:
            X = self.conv3(X)
        Y += X
        return F.relu(Y)


class ResNet18(nn.Module):
    def __init__(self, residual_block = Residual, num_classes=1000):
        super(ResNet18, self).__init__()
        self.b1 = nn.Sequential(nn.Conv2d(in_channels = 3, out_channels = 64, 
                                    kernel_size=7, stride=2, padding=3),
                        nn.BatchNorm2d(64), nn.ReLU(),
                        nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
        input_channels = 64
        num_channels = 64 
        self.b23 = nn.Sequential(residual_block(num_channels, num_channels))
        self.b45 = nn.Sequential(residual_block(num_channels, num_channels))
        input_channels = 64
        num_channels = 128
        self.b67 = nn.Sequential(residual_block(input_channels, num_channels, 
                                                use_1x1conv=True, strides=2))
        self.b89 = nn.Sequential(residual_block(num_channels, num_channels))
        input_channels = 128
        num_channels = 256
        self.b1011 = nn.Sequential(residual_block(input_channels, num_channels, 
                                                  use_1x1conv=True, strides=2))
        self.b1213 = nn.Sequential(residual_block(num_channels, num_channels))
        input_channels = 256
        num_channels = 512
        self.b1415 = nn.Sequential(residual_block(input_channels, num_channels, 
                                                  use_1x1conv=True, strides=2))
        self.b1617 = nn.Sequential(residual_block(num_channels, num_channels))
        self.b18 = nn.Sequential(nn.AdaptiveAvgPool2d((1,1)), nn.Flatten(), 
                                 nn.Linear(512, num_classes))   

    def forward(self, x):
        x = self.b1(x)
        x = self.b23(x)
        x = self.b45(x)
        x = self.b67(x)
        x = self.b89(x)
        x = self.b1011(x)
        x = self.b1213(x)
        x = self.b1415(x)
        x = self.b1617(x)
        x = self.b18(x)
        return x
        
    
resnet18 = ResNet18(residual_block = Residual, num_classes=7)
# net = nn.Sequential(resnet18.b1, resnet18.b23, resnet18.b45, 
#                     resnet18.b67, resnet18.b89, 
#                     resnet18.b1011, resnet18.b1213, 
#                     resnet18.b1415, resnet18.b1617, resnet18.b18) 
net = nn.Sequential(resnet18.b1, resnet18.b23, resnet18.b45, 
                    resnet18.b67, resnet18.b89, 
                    resnet18.b1011, resnet18.b1213, 
                    resnet18.b1415, resnet18.b1617, 
                    nn.AdaptiveAvgPool2d((1,1)), nn.Flatten(), nn.Linear(512, 7))    
# 输入的大小=(批量大小, 通道数, 行数, 列数)
X = torch.rand(size=(31, 3, 224, 224))
# 打印网络的大小
# 步长为2的话，高和宽减半
for layer in net:
    X = layer(X)
    print(layer.__class__.__name__,'output shape: ', X.shape) 
